# ResNet

import torch
import torchvision
import torchvision.transforms as tform
import torch.nn as nn
import torch.optim as optim
from torchvision import models

datadir = "/workspace/small/train"
torch.manual_seed(5)

transform = tform.Compose([
    tform.Resize(256),
    tform.CenterCrop(224),
    tform.ToTensor(),
    tform.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
    ])

trainset = torchvision.datasets.ImageFolder(root=datadir, transform=transform)
trainloader = torch.utils.data.DataLoader(trainset, batch_size=512, shuffle=True, num_workers=4)

model = models.resnet18(pretrained=True)
num_ftrs = model.fc.in_features
model.fc = nn.Linear(num_ftrs, len(trainset.classes))

criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(model.parameters(), lr=0.001, momentum=0.9)

print("Start Training")
for epoch in range(5):
    print("Start Training %d. enum=%d" % (epoch, len(trainloader)))
    running_loss = 0.0
    for i, data in enumerate(trainloader, 0):
        print("--Start Training %d, i=%d" % (epoch, i))
        inputs, labels = data
        optimizer.zero_grad()
        outputs = model(inputs)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()
        running_loss += loss.item()
        if i%100 == 99:
            print('[%d, %5d] loss: %3.f'% (epoch + 1, i+1, running_loss/100))
            running_loss = 0.0

print("Finished Training")
 
